# pip install -U langgraph
# Literal 限定变量可取的值
from typing import Literal

# dotenv 用于从 .env 文件中加载环境变量，比如 API_KEY
from dotenv import load_dotenv
# LangChain 消息类，用于封装用户消息
from langchain_core.messages import HumanMessage
# LangChain 工具装饰器，用于定义可被 LLM 调用的函数
from langchain_core.tools import tool
# OpenAI/Chat 模型接口
from langchain_openai import ChatOpenAI
# langgraph 检查点，用于在图运行中持久化状态
from langgraph.checkpoint.memory import MemorySaver
# 状态图相关类：END 表示流程结束，StateGraph 管理节点和边，MessagesState 用作状态数据载体
from langgraph.graph import END, StateGraph, MessagesState
# langgraph 预构建工具节点，用于将工具函数包装成图节点
from langgraph.prebuilt import ToolNode

# 自动加载 .env 文件中的环境变量
load_dotenv()


## -------------------------------
## 定义工具函数
## -------------------------------
@tool
def weather_tool(query: str):
    """
    模拟一个天气查询工具
    1. query 是用户输入的文本
    2. 根据 query 判断是否包含"上海"，返回不同天气信息
    """
    if "上海" in query.lower() or "shanghai" in query.lower():
        return "现在11度，有雾"
    return "现在30度，阳光明媚"


# 将工具函数放入工具列表
tools = [weather_tool]

# 将工具函数包装成 LangGraph 的 ToolNode 节点
tool_node = ToolNode(tools)

## -------------------------------
## 初始化大模型并绑定工具
## -------------------------------
llm = ChatOpenAI(
    # model="Qwen/Qwen3-14B",
     model="Qwen/Qwen2.5-VL-7B-Instruct",
    #model="xdeepseekv32exp",  # 讯飞星火
    streaming=True,  # 开启流式输出（边生成边返回）
    temperature=0.7  # 控制生成文本的随机性
).bind_tools(tools)  # 将工具绑定到模型，使模型可以自动调用


## -------------------------------
## 定义条件判断函数
## -------------------------------
def should_continue(state: MessagesState) -> Literal["tools", END]:
    """
       决定下一个节点调用哪个节点
       1. state['messages'] 存放所有对话消息
       2. last_message.tool_calls 判断 LLM 是否调用了工具
       3. 如果调用了工具，则转到 "tools" 节点
       4. 否则流程结束（END）
       """
    messages = state['messages']
    last_message = messages[-1]
    # 如果LLM 调用了工具，则转到“tools”节点
    # 只有 模型调用工具（function calling）时，对应消息才会有 tool_calls。 它记录了模型调用的 工具名称、参数、调用时间等信息。  如果模型没有调用工具，这个字段通常为空或 None。
    if last_message.tool_calls:
        return "tools"
    return END  # 否则，停止（回复客户）


## -------------------------------
## 定义调用模型的函数
## -------------------------------
def call_model(state: MessagesState):
    """
       调用大模型生成回复
       1. state['messages'] 是当前对话的消息列表
       2. llm.invoke(messages) 调用 LLM 获取响应
       3. 返回 {"messages": [response]} 形式，用于更新状态
       """
    messages = state['messages']
    response = llm.invoke(messages)
    return {"messages": [response]}  # 返回列表形式，以便状态更新


## -------------------------------
## 构建状态图
## -------------------------------
workflow = StateGraph(MessagesState)  # 使用 MessagesState 作为状态载体

# 添加节点
workflow.add_node("agent", call_model)  # agent 节点，调用 LLM
workflow.add_node("tools", tool_node)  # tools 节点，调用外部工具

# 设置入口节点，第一个执行的节点
workflow.set_entry_point("agent")

# 条件边：从 agent 根据 should_continue 决定下一个节点
workflow.add_conditional_edges(
    # 首先定义其实节点，使用agent
    "agent",
    # 然后 传递决定下一个调用节点的函数
    should_continue
)

# 普通边：从 tools 节点完成后回到 agent 节点
workflow.add_edge("tools", "agent")

## -------------------------------
## 持久化状态
## -------------------------------
check_pointer = MemorySaver()  # 内存检查点，用于在图运行之间保存状态

## -------------------------------
## 编译图
## -------------------------------
# 将状态图编译成可执行对象 app，可像普通函数一样调用
app = workflow.compile(checkpointer=check_pointer)

## -------------------------------
## 执行图
## -------------------------------
final_state = app.invoke(
    {"messages": [HumanMessage(content="日本的天气怎么样？")]},  # 初始消息
    config={"configurable": {"thread_id": 517}}  # 可配置参数，例如线程 ID
)

# 获取最终 LLM 回复内容 [-1] 表示 取最后一条消息
result = final_state["messages"][-1].content
print(result)

# 从 final_state 中获取最后一条消息的内容
# final_state = app.invoke(
#     {"messages": [HumanMessage(content="我问的是哪个城市？")]},  # 初始消息
#     config={"configurable": {"thread_id": 517}}  # 可配置参数，例如线程 ID
# )
# result = final_state["messages"][-1].content
# print(result)

## -------------------------------
## 可视化状态图
## -------------------------------
graph_png = app.get_graph().draw_mermaid_png()
with open("llm_langgraph_0_入门.png", "wb") as f:
    f.write(graph_png)
